# -*- coding: utf-8 -*-
from flask import Flask, render_template, request, jsonify
import json
import os
from pathlib import Path
from openai import OpenAI
import re
from auth import auth_login, require_login
import requests

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = './'
app.config['ALLOWED_EXTENSIONS'] = {'docx', 'doc', 'pdf'}

client = OpenAI(
    # 若没有配置环境变量，请用百炼API Key将下行替换为：api_key="sk-xxx",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)


def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
        filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


@app.route('/login', methods=['POST', 'GET'])
def login():
    return auth_login()


@app.route('/', methods=['GET'])
@require_login
def index():
    return render_template('index.html')


@app.route('/api/upload', methods=['POST'])
@require_login
def handle_upload():
    """上传文档文件接口"""
    try:
        url = request.args.get('url')
        file = request.files.get('file')
        if not url and not file:
            return jsonify({"error": "file part must be provided"}), 400
        if url:
            filename = url.split('/')[-1]
            if not allowed_file(filename):
                return jsonify({"error": "File type not allowed"}), 400
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            if not download_from_oss(url, file_path):
                return jsonify({"error": "No selected file"}), 400
        elif 'file' in request.files:
            file = request.files['file']
            if file.filename == '':
                return jsonify({"error": "No selected file"}), 400
            if not allowed_file(file.filename):
                return jsonify({"error": "File type not allowed"}), 400
            filename = file.filename
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(file_path)
        else:
            return jsonify({"error": "No file part"}), 400

        file_object = client.files.create(file=Path(file_path), purpose="file-extract")
        os.remove(file_path)
        data = {'file_id': file_object.id, 'filename': file_object.filename}
        return jsonify({'code': 200, 'data': data})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/api/completions', methods=['POST'])
@require_login
def handle_completions():
    """信息提取接口"""
    try:
        data = request.json
        file_ids = data.get('file_ids')
        prompt = data.get('prompt', '')
        resp_content = completion_request(file_ids, prompt)
        resp_content = json.loads(resp_content)
        resp_content = resp_content['choices'][0]['message']['content'];
        resp_content = extract_code_blocks(resp_content)
        if resp_content is None:
            return jsonify({'code': 400, 'data': {}})
        data = json.loads(resp_content)
        return jsonify({'code': 200, 'data': data, 'file_ids': file_ids})
    except Exception as e:
        return jsonify({'error': str(e)}), 500


def download_from_oss(url, destination):
    """下载文档"""
    response = requests.get(url)
    if response.status_code == 200:
        with open(destination, 'wb') as f:
            f.write(response.content)
        return True
    return False


def completion_request(file_ids, extra_info):
    """调用百炼模型服务"""
    # 详细API使用请参见 https://help.aliyun.com/zh/model-studio/developer-reference/openai-file-interface
    # 任务异步执行请参见 https://help.aliyun.com/zh/model-studio/developer-reference/batch-interfaces-compatible-with-openai
    prompt = (
        "# 角色\n"
        "你是一名文档文件分析师，擅长对各种文档内容进行提取。\n\n"
        "# 任务描述\n"
        "给你一几个文档文件，请你完成以下任务。\n"
        "- 分析文档内容，提取关键信息，总结摘要。\n"
        "# 限制\n"
        "- 总结时需严格依据文档内容，不可添加个人臆测或创意性内容。\n"
        "- 只能使用JSON返回结果，key必须使用中文，JSON数据格式必须正确，百分数用字符串返回。例如：0.1%\n"
        "# 输入数据\n"
        "## 文档提取内容补充信息 (你需要针对补充信息，进行文档内容提取，如果输入为空则忽略补充信息)\n"
        f"{extra_info}\n\n"
        "# 输出格式\n"
        "直接JSON格式输出内容。\n"
    )
    messages = []
    for file_id in file_ids:
        item = {'role': 'system', 'content': f'fileid://' + file_id}
        messages.append(item)
    messages.append({'role': 'user', 'content': prompt})
    completion = client.chat.completions.create(
        model="qwen-long",
        messages=messages
    )
    return completion.model_dump_json()


def extract_code_blocks(text):
    """ 正则表达式匹配  json包裹的代码块 """
    pattern = r'```json\n(.*?)```'
    # 使用re.findall()函数查找所有匹配的代码块
    code_blocks = re.findall(pattern, text, re.DOTALL)
    if len(code_blocks) > 0:
        return code_blocks[0]
    return None


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=9000)
